/*-----------------------------------------------------------------------------
 * file:  sr_ip.h 
 * date:  Thu Oct 18, 2012  
 * Author: yujiali@cs.toronto.edu
 *
 * Description:
 *
 * Methods and datastructures for handeling the ARP packets 
 *---------------------------------------------------------------------------*/

#include <assert.h>
#include <stdio.h>
#include <arpa/inet.h>
#include <sys/types.h>
#include <string.h>

#include "sr_router.h"
#include "sr_protocol.h"
#include "sr_icmp.h"
#include "sr_if.h"
#include "sr_rt.h"
#include "sr_ip.h"
#include "sr_arp.h"

/* shared packet buffer */
uint8_t ip_pkt_buf[IP_MAXPACKET];
uint8_t ip_pkt_hdr_buf[IP_HDR_MAX_LEN];

int sr_send_ip_to_ethernet(struct sr_instance*, struct ip*, uint32_t, 
        uint8_t*, unsigned int, struct sr_if*);
int sr_forward_ip_packet(struct sr_instance*, struct ip*, 
        uint8_t*, unsigned int);
int sr_get_rtentry_iface(struct sr_instance* sr, uint32_t ip_dst, 
        struct sr_rt** rt_entry, struct sr_if** iface);

/* only for debugging */
/* void sr_print_iphdr(struct ip* ip_hdr); */


/*--------------------------------------------------------------------- 
 * Method: sr_handle_ip(struct sr_instance* sr,uint8_t* packet, 
 *      unsigned int len, char* interface)
 *
 * Scope: local
 *
 * This method is called every time an IP packet is received. The packet
 * buffer is complete with Ethernet headers and IP packet is inside it.
 *
 * After processing the IP packet, a 0 is returned if everything goes
 * well, otherwise an error code is returned.
 *
 *---------------------------------------------------------------------*/
int sr_handle_ip(struct sr_instance* sr, 
        uint8_t* packet, 
        unsigned int len, 
        char* interface)
{
    /* REQUIRES */
    assert(sr);
    assert(packet);
    assert(interface);

    struct ip* ip_hdr = (struct ip *)(packet + sizeof(struct sr_ethernet_hdr));
    struct sr_if* iface = sr_get_interface(sr, interface);
    uint32_t ip_dst;

    assert(iface);

    /* sr_print_iphdr(ip_hdr); */

    /* perform a few checks */

    if (ip_hdr->ip_v != IP_V4)
    {
        fprintf(stderr, "Not an IPv4 packet, version: %u\n", ip_hdr->ip_v);
        return -1;
    }

    if (ip_hdr->ip_hl < 5)
    {
        fprintf(stderr, "IP header size too small, provided: %u\n", 
                ip_hdr->ip_hl);
        return -1;
    }
    

    if (sr_checksum((uint16_t*)ip_hdr, ((unsigned int)(ip_hdr->ip_hl)) << 1,
            CHECKSUM_DONT_PAD))
    {
        fprintf(stderr, "IP checksum failed...\n");
        return -1;
    }

    /* check if this packet's destination is the current node */
    ip_dst = ip_addr_to_uint32(ip_hdr->ip_dst);

    /* destination is not the current node, should route it correctly */
    if (!sr_search_if_ip(sr, ip_dst))
    {
        /* send error message back if TTL reaches 0 */
        if (ip_hdr->ip_ttl - 1 == 0)
        {
            fprintf(stderr, "IP TTL reaches 0...\n");
            sr_send_icmp_packet(sr, 
                    ICMP_TYPE_TIME_EXCEEDED, 
                    ICMP_CODE_TIME_EXCEEDED_IN_TRANSIT,
                    ip_addr_to_uint32(ip_hdr->ip_src), 
                    iface->ip, 
                    (uint8_t *)ip_hdr);
            return -1;
        }

        return sr_forward_ip_packet(sr, ip_hdr, 
                packet + sizeof(struct sr_ethernet_hdr) + sizeof(struct ip), 
                len - sizeof(struct sr_ethernet_hdr) - sizeof(struct ip));
    }
    
    /* handle IP packets according to different upper layer protocols */
    switch(ip_hdr->ip_p)
    {
        case IPPROTO_ICMP:
            printf("Received ICMP packet...\n");
            sr_handle_icmp(sr, 
                    packet + sizeof(struct sr_ethernet_hdr),
                    len - sizeof(struct sr_ethernet_hdr));
            break;
        case IPPROTO_TCP:
        case IPPROTO_UDP:
            printf("Received Transport layer packet...\n");

            sr_send_icmp_packet(sr, ICMP_TYPE_DST_UNREACHABLE, 
                    ICMP_CODE_PORT_UNREACHABLE, 
                    ip_addr_to_uint32(ip_hdr->ip_src), ip_dst, (uint8_t *)ip_hdr);
            break;
        default:
            printf("Unknown packet type: 0x%02d\n", ip_hdr->ip_p);
    }

    return 0;
}

/*--------------------------------------------------------------------- 
 * Method: sr_send_ip_to_ethernet(...)
 *
 * Scope: local
 *
 * This method is called every time an IP packet is to be resent. The sender
 * should prepare an IP header complete with all required fields and options
 * before calling this function. This function would not change the header 
 * and data.
 *
 * After processing the IP packet, a 0 is returned if everything goes
 * well, otherwise an error code is returned.
 *
 *---------------------------------------------------------------------*/
int sr_send_ip_to_ethernet(struct sr_instance* sr, 
        struct ip* sender_ip_hdr,
        uint32_t gw_ip,
        uint8_t* packet, 
        unsigned int len, 
        struct sr_if* iface)
{
    assert(sr);
    assert(packet);
    assert(iface);

    uint8_t ether_addr_dst[ETHER_ADDR_LEN];

    struct sr_ethernet_hdr* e_hdr = (struct sr_ethernet_hdr *)ip_pkt_buf;
    struct ip* ip_hdr = 
        (struct ip *)(ip_pkt_buf + sizeof(struct sr_ethernet_hdr));
    unsigned int ip_hdr_len = 0;

    printf("sr_send_ip_to_ethernet\n");

    ip_hdr_len = sender_ip_hdr->ip_hl;
    ip_hdr_len = ip_hdr_len << 2;

    /* ARP entry not found, send ARP requests */
    if (!sr_search_arp(sr, gw_ip, ether_addr_dst))
    {
        sr_request_arp(sr, gw_ip, iface);
        sr_add_to_queue(sr, gw_ip, iface, (uint8_t*)sender_ip_hdr, ip_hdr_len + len);
        return -1;
    }

    memcpy(e_hdr->ether_dhost, ether_addr_dst, ETHER_ADDR_LEN);
    memcpy(e_hdr->ether_shost, iface->addr, ETHER_ADDR_LEN);
    e_hdr->ether_type = htons(ETHERTYPE_IP);

    memcpy(ip_hdr, sender_ip_hdr, ip_hdr_len);
    memcpy(ip_pkt_buf + sizeof(struct sr_ethernet_hdr) + ip_hdr_len, packet, len);

    /* sr_print_iphdr(ip_hdr); */

    /* send the packet out through VNS interface */
    return sr_send_packet(sr, ip_pkt_buf, 
            len + sizeof(struct sr_ethernet_hdr) + ip_hdr_len, 
            iface->name);
}

/*--------------------------------------------------------------------- 
 * Method: sr_send_ip_to_ethernet_thread_safe(...)
 *
 * Scope: global
 *
 * This method send complete IP packets to the ethernet. The buffer is large
 * enough to contain the whole ethernet header (though they are left empty).
 * Since now the buffer is now not a shared resource, this function is thread
 * safe.
 *
 * buf is a packet buffer, starts with ethernet headers followed by a
 * complete IP packet.
 *---------------------------------------------------------------------*/
int sr_send_ip_to_ethernet_thread_safe(struct sr_instance* sr, uint8_t* buf)
{
    assert(sr);
    assert(buf);

    struct sr_ethernet_hdr* e_hdr = (struct sr_ethernet_hdr *)buf;
    struct ip* ip_hdr = (struct ip *)(buf + sizeof(struct sr_ethernet_hdr));
    struct sr_rt* rt_entry = 0;
    struct sr_if* iface = 0;
    uint32_t ip_gw = 0;

    /* look up in the routing table */
    if (sr_get_rtentry_iface(sr, ip_addr_to_uint32(ip_hdr->ip_dst),
            &rt_entry, &iface))
        return -1;

    ip_gw = ip_addr_to_uint32(rt_entry->gw);

    /* ARP entry not found, send ARP requests */
    if (!sr_search_arp(sr, ip_gw, e_hdr->ether_dhost))
    {
        sr_request_arp(sr, ip_gw, iface);
        sr_add_to_queue(sr, ip_gw, iface, (uint8_t*)ip_hdr, ntohs(ip_hdr->ip_len));
        return -1;
    }

    memcpy(e_hdr->ether_shost, iface->addr, ETHER_ADDR_LEN);
    e_hdr->ether_type = htons(ETHERTYPE_IP);

    return sr_send_packet(sr, buf, 
            sizeof(struct sr_ethernet_hdr) + ntohs(ip_hdr->ip_len),
            iface->name);
    return 0;
}


/*--------------------------------------------------------------------- 
 * Method: sr_forward_ip_packet(...)
 *
 * Scope: local
 *
 * This method is used to forward IP packets, used in routing. Other than
 * decrementing TTL and updating checksum, the original IP header is not
 * changed.
 *
 * After processing the IP packet, a 0 is returned if everything goes
 * well, otherwise an error code is returned.
 *
 *---------------------------------------------------------------------*/
int sr_forward_ip_packet(struct sr_instance* sr, struct ip* sender_ip_hdr, 
        uint8_t* buf, unsigned int len)
{
    assert(sr);
    assert(buf);

    struct sr_rt* rt_entry = 0;
    struct sr_if* iface = 0;

    printf("sr_forward_ip_packet\n");

    if (sr_get_rtentry_iface(sr, ip_addr_to_uint32(sender_ip_hdr->ip_dst),
            &rt_entry, &iface))
        return -1;

    /* update TTL and checksum */
    sender_ip_hdr->ip_ttl -= 1;
    sender_ip_hdr->ip_sum = 0;
    sender_ip_hdr->ip_sum = sr_checksum(
            (uint16_t*)sender_ip_hdr, (sender_ip_hdr->ip_hl) << 1, 
            CHECKSUM_DONT_PAD);

    /* send the packet down to Ethernet layer */
    return sr_send_ip_to_ethernet(sr, sender_ip_hdr, 
            ip_addr_to_uint32(rt_entry->gw), buf, len, iface);
}

/*--------------------------------------------------------------------- 
 * Method: sr_send_ip_pkt_w_hdr(...)
 *
 * Scope: global
 *
 * This method is used to send IP packets, with the IP header prepared
 * by the user. Nothing would be changed for the IP header and data.
 * It is almost identical to sr_forward_ip_packet, except that TTL and
 * checksum are left unchanged.
 *
 * After processing the IP packet, a 0 is returned if everything goes
 * well, otherwise an error code is returned.
 *---------------------------------------------------------------------*/

int sr_send_ip_pkt_w_hdr(struct sr_instance* sr, struct ip* ip_hdr, 
        uint8_t* buf, unsigned int len)
{
    assert(sr);
    assert(buf);

    struct sr_rt* rt_entry = 0;
    struct sr_if* iface = 0;

    printf("sr_send_ip_pkt_w_hdr\n");

    if (sr_get_rtentry_iface(sr, ip_addr_to_uint32(ip_hdr->ip_dst),
            &rt_entry, &iface))
        return -1;

    /* send the packet down to Ethernet layer */
    return sr_send_ip_to_ethernet(sr, ip_hdr, 
            ip_addr_to_uint32(rt_entry->gw), buf, len, iface);

}

/*--------------------------------------------------------------------- 
 * Method: sr_send_ip_packet(...)
 *
 * Scope: global
 *
 * This method is used by upper layers to send IP packets. The current
 * implementation does not support fragmentation, so the user of this
 * function should make sure that the length of the packet sent should
 * be small enough to fit to the local network.
 *
 * TODO: handle packet fragmentation
 * 
 * After processing the IP packet, a 0 is returned if everything goes
 * well, otherwise an error code is returned.
 *
 *---------------------------------------------------------------------*/
int sr_send_ip_packet(struct sr_instance* sr, uint32_t ip_src, uint32_t ip_dst,
        uint8_t prot, uint8_t tos, uint8_t ttl, uint8_t* buf, unsigned int len,
        uint16_t id, uint8_t df, uint8_t* opt, unsigned int opt_len)
{
    assert(sr);
    assert(buf);

    struct ip* ip_hdr = (struct ip *)ip_pkt_hdr_buf;
    struct sr_rt* rt_entry = 0;
    struct sr_if* iface = 0;
    
    /* search the routing table */
    if (sr_get_rtentry_iface(sr, ip_dst, &rt_entry, &iface))
        return -1;
    
    /* assemble an IP header */
    ip_hdr->ip_hl = IP_HDR_MIN_WORDS + (opt_len >> 2);
    ip_hdr->ip_v = IP_V4;
    ip_hdr->ip_tos = tos;
    ip_hdr->ip_len = htons(IP_HDR_MIN_LEN + opt_len + len);
    ip_hdr->ip_id = htons(id);
    if (df)
        ip_hdr->ip_off = htons(IP_DF);
    else
        ip_hdr->ip_off = 0;

    ip_hdr->ip_ttl = ttl;
    ip_hdr->ip_p = prot;
    ip_hdr->ip_sum = 0;
    ip_hdr->ip_src = ip_addr_to_in_addr(ip_src);
    ip_hdr->ip_dst = ip_addr_to_in_addr(ip_dst);

    /* handle the option part */
    if (opt_len > 0)
    {
        assert(opt);
        memcpy(((uint8_t *)ip_hdr) + sizeof(struct ip), opt, opt_len);
    }

    ip_hdr->ip_sum = sr_checksum((uint16_t *)ip_hdr, ip_hdr->ip_len, 
            CHECKSUM_DONT_PAD);

    /* send the packet to the ethernet */
    return sr_send_ip_to_ethernet(sr, ip_hdr, ip_addr_to_uint32(rt_entry->gw),
            buf, len, iface);
}

/*--------------------------------------------------------------------- 
 * Method: sr_get_rtentry_iface(...)
 *
 * Scope: global
 *
 * A wrapper function to perform a routing table search and an interface
 * list search.
 * 
 * After processing the request, a 0 is returned if everything goes
 * well, otherwise an error code is returned.
 *
 *---------------------------------------------------------------------*/
int sr_get_rtentry_iface(struct sr_instance* sr, uint32_t ip_dst, 
        struct sr_rt** rt_entry, struct sr_if** iface)
{
    /* look up for the destination IP in the routing table */
    *rt_entry = sr_search_rt_ip(sr, ip_dst);
    if (!(*rt_entry))
    {
        fprintf(stderr, "Destination IP not found in routing table.\n");
        return -1;
    }

    /* search for the required interface */
    *iface = sr_get_interface(sr, (*rt_entry)->interface);
    if (!(*iface))
    {
        fprintf(stderr, "Cannot find the required interface: %s\n",
                (*rt_entry)->interface);
        return -1;
    }
    return 0;
}

/*--------------------------------------------------------------------- 
 * Method: sr_checksum(...)
 *
 * Scope: global
 *
 * This method computes the a 16-bit checksum of a stream of 16-bit data.
 * Can be used for IP and layers above. The data stream can be padded with
 * one extra byte of zeros.
 *
 *---------------------------------------------------------------------*/
uint16_t sr_checksum(uint16_t* buf, unsigned int count, unsigned int pad)
{
    uint32_t sum = 0;

    if (pad == CHECKSUM_DONT_PAD)
    {
        while (count--)
        {
            sum += *buf;
            buf++;

            if (sum & 0xffff0000)
            {
                sum &= 0xffff;
                sum++;
            }
        }
        return ~(sum & 0xffff);
    }
    else
    {
        while (count > 1)
        {
            count--;
            sum += *buf;
            buf++;

            if (sum & 0xffff0000)
            {
                sum &= 0xffff;
                sum++;
            }
        }

        /* pad the last byte with another extra byte of 0 */
        sum += (*buf) & 0xff00;
        if (sum & 0xffff0000)
        {
            sum &= 0xffff;
            sum++;
        }

        return ~(sum & 0xffff);
    }
}



/* only for debugging */
void sr_print_iphdr(struct ip* ip_hdr)
{
    assert(ip_hdr);

    char ip_addr[256];

    printf("ip_v:   %u\n", ip_hdr->ip_v);
    printf("ip_hl:  %u\n", ip_hdr->ip_hl);
    printf("ip_tos: %u\n", ip_hdr->ip_tos);
    printf("ip_len: %u\n", ntohs(ip_hdr->ip_len));
    printf("ip_id:  %u\n", ntohs(ip_hdr->ip_id));
    printf("ip_rf:  %u\n", (ntohs(ip_hdr->ip_off) & IP_RF) >> 15);
    printf("ip_df:  %u\n", (ntohs(ip_hdr->ip_off) & IP_DF) >> 14);
    printf("ip_mf:  %u\n", (ntohs(ip_hdr->ip_off) & IP_MF) >> 13);
    printf("ip_off: %u\n", ntohs(ip_hdr->ip_off) & IP_OFFMASK);
    printf("ip_ttl: %u\n", ip_hdr->ip_ttl);
    printf("ip_p:   0x%x\n", ip_hdr->ip_p);
    printf("ip_sum: %u\n", ip_hdr->ip_sum);
    
    inet_ntop(AF_INET, &(ip_hdr->ip_src), ip_addr, 256);
    printf("ip_src: %s\n", ip_addr);
    inet_ntop(AF_INET, &(ip_hdr->ip_dst), ip_addr, 256);
    printf("ip_dst: %s\n", ip_addr);
}

